package org.yinweichen.llm_dataset_backend.entity;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.persistence.*;
import lombok.*;
import org.hibernate.annotations.JdbcTypeCode;
import org.hibernate.type.SqlTypes;
import org.springframework.data.annotation.LastModifiedDate;

import java.io.IOException;
import java.util.Date;
import java.util.List;

@Data
@Entity
@Table(name = "eval_tasks")
@NoArgsConstructor
@AllArgsConstructor
@Builder
public class EvalTask {
    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    @Column(name = "id")
    private Long id;

    @Column(name = "title")
    private String title;

    @Column(name = "description", columnDefinition = "TEXT")
    private String description;//评测任务简介

    @Column(name = "report", columnDefinition = "JSON")
    @JdbcTypeCode(SqlTypes.JSON)
    private String report;

    @Column(name = "time")
    @LastModifiedDate
    private Date time;

    @Column(name = "version")
    private String version;

//    @Column(name = "prompt", columnDefinition = "TEXT(65535)")
//    private String prompt;

    //标注评测阶段（new/answered/finished）
    @Column(name = "status")
    private String status;

    //待评测的问题列表，使用的裁判模型，待评测的模型等均保存在json
    /*
    Json格式：
    {
        prompt:"xxx",
        judges:["xxx","xxx"],
        models:["xxx","xxx"],
        questions:[0,1,2]
    }
     */
    @Column(name = "config", columnDefinition = "JSON")
    @JdbcTypeCode(SqlTypes.JSON)
    private String config;


    public String getPrompt() throws IOException {
        ObjectMapper objectMapper = new ObjectMapper();
        JsonNode rootNode = objectMapper.readTree(this.config);
        return rootNode.get("prompt").asText();
    }
    public List<String> getJudges() throws IOException {
        ObjectMapper objectMapper = new ObjectMapper();
        JsonNode rootNode = objectMapper.readTree(this.config);
        return objectMapper.convertValue(rootNode.get("judges"),
            objectMapper.getTypeFactory().constructCollectionType(List.class, String.class));
    }
    public List<String> getModels() throws IOException {
        ObjectMapper objectMapper = new ObjectMapper();
        JsonNode rootNode = objectMapper.readTree(this.config);
        return objectMapper.convertValue(rootNode.get("models"),
            objectMapper.getTypeFactory().constructCollectionType(List.class, String.class));
    }
    public List<Long> getQuestions() throws IOException {
        ObjectMapper objectMapper = new ObjectMapper();
        JsonNode rootNode = objectMapper.readTree(this.config);
        return objectMapper.convertValue(rootNode.get("questions"),
            objectMapper.getTypeFactory().constructCollectionType(List.class, Long.class));
    }
}
